package tool

import (
	"fmt"
	"sync"
	"time"
)

const (
	timestampBits = 31 // 时间戳的总位数
	machineIDBits = 1  // 机器 ID 的总位数
	sequenceBits  = 8  // 序列号的位数

	maxMachineID = -1 ^ (-1 << machineIDBits)
	sequenceMask = -1 ^ (-1 << sequenceBits)

	machineIDShift = sequenceBits
	timestampShift = sequenceBits + machineIDBits

	epoch = int64(1609459200) // 自定义基准时间戳 (2021-01-01 00:00:00 UTC)，秒级时间戳
)

type Snowflake struct {
	sync.Mutex
	timestamp int64
	machineID int64
	sequence  int64
}

func NewSnowflake(machineID int64) *Snowflake {
	if machineID > maxMachineID || machineID < 0 {
		panic(fmt.Sprintf("machine ID can't be greater than %d or less than 0", maxMachineID))
	}
	return &Snowflake{
		timestamp: 0,
		machineID: machineID,
		sequence:  0,
	}
}

func (s *Snowflake) NextID() int64 {
	s.Lock()
	defer s.Unlock()

	now := time.Now().Unix() - epoch
	if now == s.timestamp {
		s.sequence = (s.sequence + 1) & sequenceMask
		if s.sequence == 0 {
			for now <= s.timestamp {
				now = time.Now().Unix() - epoch
			}
		}
	} else {
		s.sequence = 0
	}
	s.timestamp = now

	id := (now << timestampShift) |
		(s.machineID << machineIDShift) |
		s.sequence

	return id
}
